# -*- coding:utf-8 -*-

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import os, sys
import tensorflow as tf
import time
import cv2
import argparse
import numpy as np

sys.path.append("./")
from data.io.image_preprocess import short_side_resize_for_inference_data
# sys.path.append("../libs/")
from libs.configs import cfgs
from libs.box_utils import draw_box_in_img
from help_utils import tools


def load_graph(frozen_graph_file):
    # we parse the graph_def file
    with tf.gfile.GFile(frozen_graph_file, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    # we load the graph_def in the default graph

    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def,
                            input_map=None,
                            return_elements=None,
                            name="",
                            op_dict=None,
                            producer_op_list=None)
    return graph


def inference(model, inference_save_path, real_test_imgname_list):
    graph = load_graph(model)
    input_tensor = graph.get_tensor_by_name('input_img:0')
    out_tensor = graph.get_tensor_by_name('DetResults:0')
    with tf.Session(graph=graph) as sess:
        for i, a_img_name in enumerate(real_test_imgname_list):
            raw_img = cv2.imread(a_img_name)[:, :, ::-1]
            img = raw_img - cfgs.PIXEL_MEAN
            start = time.time()
            DetResults = sess.run(out_tensor, feed_dict={input_tensor: raw_img})
            print('DetResults_shape:{}'.format(DetResults.shape))
            detected_boxes = DetResults[:, 2:]
            detected_scores = DetResults[:, 1]
            detected_categories = DetResults[:, 0]
            end = time.time()
            print("detected_boxes:{}\ndetected_scores:{}\ndetected_categories:{}".format(detected_boxes.shape,
                                                                                         detected_scores.shape,
                                                                                         detected_categories.shape))

            show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
            show_scores = detected_scores[show_indices]
            show_boxes = detected_boxes[show_indices]
            show_categories = detected_categories[show_indices]
            final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(img,
                                                                                boxes=show_boxes,
                                                                                labels=show_categories,
                                                                                scores=show_scores)
            nake_name = a_img_name.split('/')[-1]
            # print (inference_save_path + '/' + nake_name)
            cv2.imwrite(inference_save_path + '/' + nake_name,
                        final_detections[:, :, ::-1])
            print('final_detections\ncategory:{}\nscore:{}\nbox:{}'.format(show_categories, show_scores, show_boxes))
            tools.view_bar('{} image cost {}s'.format(a_img_name, (end - start)), i + 1, len(real_test_imgname_list))


def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='TestImgs...U need provide the test dir')
    parser.add_argument('--data_dir', dest='data_dir', help='data path', default='demos', type=str)
    parser.add_argument('--save_dir', dest='save_dir', help='demo imgs to save', default='inference_results', type=str)
    parser.add_argument('--GPU', dest='GPU', help='gpu id ', default='0', type=str)

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()

    return args


if __name__ == '__main__':
    args = parse_args()
    print('Called with args:')
    print(args)
    os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU
    # inference(args.data_dir, inference_save_path=args.save_dir)
    test_imgname_list = [os.path.join(args.data_dir, img_name) for img_name in os.listdir(args.data_dir)
                         if img_name.endswith(('.jpg', '.png', '.jpeg', '.tif', '.tiff'))]
    assert len(test_imgname_list) != 0, 'test_dir has no imgs there.' \
                                        ' Note that, we only support img format of (.jpg, .png, and .tiff) '
    model = './output/Pbs/model.graphdef'
    # print(model)
    inference(model, inference_save_path=args.save_dir, real_test_imgname_list=test_imgname_list)
